package org.usmile.algorithms.leetcode.middle;

import java.util.Arrays;

/**
 * 498. 对角线遍历
 *
 * 给你一个大小为 m x n 的矩阵 mat ，请以对角线遍历的顺序，用一个数组返回这个矩阵中的所有元素。
 *
 * 示例 1：
 *
 * 输入：mat = [[1,2,3],[4,5,6],[7,8,9]]
 * 输出：[1,2,4,7,5,3,6,8,9]
 * 示例 2：
 *
 * 输入：mat = [[1,2],[3,4]]
 * 输出：[1,2,3,4]
 *
 * 提示：
 *
 * m == mat.length
 * n == mat[i].length
 * 1 <= m, n <= 104
 * 1 <= m * n <= 104
 * -105 <= mat[i][j] <= 105
 */
public class _0498 {
    public static void main(String[] args) {
        int[][] mat = new int[][]{{1, 2}, {3, 4}};
        System.out.println(Arrays.toString(new _0498_Solution().findDiagonalOrder(mat)));
    }
}

class _0498_Solution {
    public int[] findDiagonalOrder(int[][] mat) {
        int rows = mat.length;
        int cols = mat[0].length;
        int[][] dirs = {{-1, 1}, {1, -1}};

        int row = 0;
        int col = 0;
        int di = 0;
        int[] result = new int[rows * cols];
        for (int i = 0; i < rows * cols; i++) {
            result[i] = mat[row][col];

            row = row + dirs[di][0];
            col = col + dirs[di][1];

            if (col >= cols) {
                col = cols - 1;
                row += 2;
                di = 1 - di;
            }
            if (row >= rows) {
                row = rows - 1;
                col += 2;
                di = 1 - di;
            }
            if (row < 0) {
                row = 0;
                di = 1 - di;
            }
            if (col < 0) {
                col = 0;
                di = 1 - di;
            }
        }

        return result;
    }
}
